from typing import TypedDict, Annotated

from langgraph.graph import add_messages, StateGraph

from common.QwenLLm import QwenLLM

llm = QwenLLM()

class State(TypedDict):
    messages:Annotated[list,add_messages]



def chatbot(state:State):
    resp = llm.invoke(state["messages"])
    return {"messages":[resp]}

grapher_builder = StateGraph(State)
grapher_builder.add_node("chat_bot",chatbot)
grapher_builder.set_entry_point("chat_bot")
graph = grapher_builder.compile()
while True:
    user_input = input("User: ")
    if user_input == "quit":
        print("good bye")
        break
    for event in graph.stream({"messages":("user",user_input)}):
        for value in event.values():
            print("assistant",value["messages"][-1].content)
